# -*- coding: utf-8 -*-

# ImportError: /home/kzl/anaconda2/bin/../lib/libstdc++.so.6: version `GLIBCXX_3.4.21' not found
# https://blog.csdn.net/a8039974/article/details/78544209
from __future__ import print_function
from coviar import load
from coviar import get_num_frames
import numpy as np
from transforms import color_aug
from tqdm import tqdm
import pandas as pd
import scipy.io as sio
import os
import h5py  # 导入工具包
import argparse

parser = argparse.ArgumentParser(description="gen tvsum/ovp/summme/youtube summary frame data")

parser.add_argument('--ft', type=str, choices=['i', 'm', 'r'], help='dataset name.')
args = parser.parse_args()

GOP_SIZE = 12

def clip_and_scale(img, size):
    return (img * (127.5 / size)).astype(np.int32)


if __name__ == "__main__":
    top = '/home/ping501b/data/videosum/videosum_concise_raw'
    # out = '/home/ping501b/data/vdo_summary_coviar_raw'

    h5f = h5py.File(f'vdo_summary_{args.ft}.h5','w')  # 创建一个h5文件，文件指针是h5f
    # http://docs.h5py.org/en/latest/high/group.html#Group
    for dr in os.listdir(top):  # ovp, tvsum, summe, youtube
        h5f.create_group(dr)
        # for fn in tqdm(os.listdir(f'{top}/{dr}')):
        for fn in os.listdir(f'{top}/{dr}'):
            path = f'{top}/{dr}/{fn}'
            num_frms = get_num_frames(path)
            # h5f[dr][fn] = []
            # h5f[dr].create_group(fn)
            frms = []
            # for i in tqdm(range(num_frms)):
            for i in range(num_frms):
                if args.ft == 'i':
                    if i % 12 == 0:
                        ifrm = load(path, int(i/12), 0, 0, True)
                        frms.append(ifrm[..., ::-1])  # 通道维反向
                elif args.ft == 'm':
                    if i % 12 != 0:
                        mv = load(path, int(i/12), i%12, 1, True)
                        mv = clip_and_scale(mv, 20)
                        mv += 128  # TODO 运动尺度的均值
                        # <0 -> =0 && >255 -> =255
                        mv = (np.minimum(np.maximum(mv, 0), 255)).astype(np.uint8)
                        frms.append(mv)
                elif args.ft == 'r':
                    if i % 12 != 0:
                        res = load(path, int(i/12), i%12, 2, True)
                        res += 128
                        res = (np.minimum(np.maximum(res, 0), 255)).astype(np.uint8)
                        frms.append(res)
            tmp = np.asarray(frms)
            dset = h5f[f'/{dr}'].create_dataset(fn, data=tmp)

    h5f.close()
